package com.cd.university.aop;

import com.cd.university.common.RestResult;
import com.cd.university.common.UniversityConstants;
import net.jodah.expiringmap.ExpirationPolicy;
import net.jodah.expiringmap.ExpiringMap;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.web.authentication.logout.SecurityContextLogoutHandler;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @author 陈栋
 * @create 2021/12/13 15:32
 */
@Aspect
@Component
public class LimitRequestAspect {

    /**
     * key为访问的url，而ExpiringMap中的key为ip地址，value为访问次数
     */
    private static ConcurrentHashMap<String, ExpiringMap<String, Integer>> requestLimitMap = new ConcurrentHashMap<>();

    @Pointcut("execution(* com.cd.university.controller..*.*(..))")
    public void limitRequest() {

    }

    @Around("execution(* com.cd.university.controller..*.*(..))")
    public Object doAround(ProceedingJoinPoint pjp) throws Throwable {
        // 获取RequestAttributes
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        // 从获取RequestAttributes中获取HttpServletRequest的信息
        HttpServletRequest request = (HttpServletRequest) requestAttributes
                .resolveReference(RequestAttributes.REFERENCE_REQUEST);
        // 从获取RequestAttributes中获取HttpServletResponse的信息
        HttpServletResponse response =((ServletRequestAttributes)RequestContextHolder.getRequestAttributes()).getResponse();

        String requestURI = request.getRequestURI();
        // 如果没有就设置一个新创建的ExpiringMap，否则就取出原来的ExpiringMap
        ExpiringMap<String, Integer> uc = requestLimitMap.getOrDefault(requestURI, ExpiringMap.builder().variableExpiration().build());
        // 取出改ip地址访问该接口的次数，如果没有默认设置为0
        String addr = request.getRemoteAddr();
        Integer uCount = uc.getOrDefault(addr, 0);

        // 如果在ExpiringMap过期之前，请求次数超过了默认值5，那么就直接抛出异常
        if (uCount >= UniversityConstants.MAX_LIMIT_REQUEST_COUNT) {
            // 退出当前用户
            Authentication auth = SecurityContextHolder.getContext().getAuthentication();
            new SecurityContextLogoutHandler().logout(request,response,auth);
            return RestResult.failure("请求次数超过限制,休息一下吧！");
        }else if (uCount==0) {
            // 第一次请求
            uc.put(addr,uCount+1, ExpirationPolicy.CREATED,UniversityConstants.MAX_LIMIT_REQUEST_TIME, TimeUnit.MILLISECONDS);
        }else {
            // 其他情况直接+1
            uc.put(addr, uCount + 1);
        }

        requestLimitMap.put(requestURI, uc);
        // 真正执行方法
        return pjp.proceed();
    }
}
